import os
import pandas as pd

csv_path = r"dataset/valid_label.csv"
aim_path = r"dataset/valid/label/"
if not os.path.exists(aim_path):
    os.mkdir(aim_path)

# 读取文件
data_file = pd.read_csv(csv_path, header=None)
# print(data_file)
# print(data_file.iloc[0,0])

# 取出名字
file_name = data_file.iloc[:, 0]
file_label = data_file.iloc[:, 1]

for file_index, loop in enumerate(file_name):
    # str2num
    label_str = file_label[file_index].split('[')[1].split(']')[0]
    if label_str == "snow":
        img_label = 0
    elif label_str == "lily":
        img_label = 1
    elif label_str == "wind":
        img_label = 2
    else:
        print("请继续增加分类！")

    with open(os.path.join(aim_path, loop.split('.')[0] + ".txt"), 'w') as fp:
        fp.write(str(img_label))